{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "027ff9dc-dd98-4846-8edd-5182a7acd2ac",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/cuongvng/opt/miniconda3/envs/ecg/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import joblib\n",
    "import h5py\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.model_selection import train_test_split\n",
    "import warnings; warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "3f042dca-4d88-44de-a7fc-28795307c07a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "CLASSES = ['NORM', 'MI']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3c0c3a27-2f0a-4630-9d3e-6f008a277ba3",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "data = np.load('./data/data_raw.npz')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "39d8e22c-f2c6-4846-8ef5-ac425f2b4371",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "KeysView(NpzFile './data/data_raw.npz' with keys: patient212/s0434_re, patient195/s0337lre, patient151/s0206_re, patient116/s0302lre, patient075/s0327lre...)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c8b9a5ab-72b8-4308-afd0-d485824f9617",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>patient</th>\n",
       "      <th>record_id</th>\n",
       "      <th>sig_name</th>\n",
       "      <th>Reason_for_admission</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>patient001</td>\n",
       "      <td>s0010_re</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>patient001</td>\n",
       "      <td>s0014lre</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>patient001</td>\n",
       "      <td>s0016lre</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>patient002</td>\n",
       "      <td>s0015lre</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>patient003</td>\n",
       "      <td>s0017lre</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>544</th>\n",
       "      <td>patient292</td>\n",
       "      <td>s0555_re</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>545</th>\n",
       "      <td>patient292</td>\n",
       "      <td>s0556_re</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>546</th>\n",
       "      <td>patient293</td>\n",
       "      <td>s0557_re</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>547</th>\n",
       "      <td>patient293</td>\n",
       "      <td>s0558_re</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>548</th>\n",
       "      <td>patient294</td>\n",
       "      <td>s0559_re</td>\n",
       "      <td>['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...</td>\n",
       "      <td>Myocardial infarction</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>549 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "        patient record_id                                           sig_name  \\\n",
       "0    patient001  s0010_re  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "1    patient001  s0014lre  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "2    patient001  s0016lre  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "3    patient002  s0015lre  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "4    patient003  s0017lre  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "..          ...       ...                                                ...   \n",
       "544  patient292  s0555_re  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "545  patient292  s0556_re  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "546  patient293  s0557_re  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "547  patient293  s0558_re  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "548  patient294  s0559_re  ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', ...   \n",
       "\n",
       "      Reason_for_admission  \n",
       "0    Myocardial infarction  \n",
       "1    Myocardial infarction  \n",
       "2    Myocardial infarction  \n",
       "3    Myocardial infarction  \n",
       "4    Myocardial infarction  \n",
       "..                     ...  \n",
       "544  Myocardial infarction  \n",
       "545  Myocardial infarction  \n",
       "546  Myocardial infarction  \n",
       "547  Myocardial infarction  \n",
       "548  Myocardial infarction  \n",
       "\n",
       "[549 rows x 4 columns]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta = pd.read_csv(\"./data/meta.csv\", usecols=[\"patient\", \"record_id\", \"sig_name\", \"Reason_for_admission\"])\n",
    "meta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "98a486c5-fefd-42ae-829f-168d2b42e934",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0      patient001/s0010_re\n",
       "1      patient001/s0014lre\n",
       "2      patient001/s0016lre\n",
       "3      patient002/s0015lre\n",
       "4      patient003/s0017lre\n",
       "              ...         \n",
       "544    patient292/s0555_re\n",
       "545    patient292/s0556_re\n",
       "546    patient293/s0557_re\n",
       "547    patient293/s0558_re\n",
       "548    patient294/s0559_re\n",
       "Name: key, Length: 549, dtype: object"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "meta['key'] = meta['patient'] + '/' + meta['record_id']\n",
    "meta['key']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "063e1ccc-093e-4bb8-aa06-faaa8a36c57d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0      [0, 1]\n",
       "1      [0, 1]\n",
       "2      [0, 1]\n",
       "3      [0, 1]\n",
       "4      [0, 1]\n",
       "        ...  \n",
       "517    [0, 1]\n",
       "518    [0, 1]\n",
       "519    [0, 1]\n",
       "520    [0, 1]\n",
       "521    [0, 1]\n",
       "Name: label, Length: 522, dtype: object"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_label(d):\n",
    "    if d == 'Myocardial infarction':\n",
    "        return np.array([0, 1])\n",
    "    elif d == 'Healthy control':\n",
    "        return np.array([1, 0])\n",
    "    else:\n",
    "        return np.array([0, 0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "864874f0-d764-4f9b-88fe-3f03bd1e4e8e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "meta = meta.dropna().reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "b0f394b3-f5a2-4515-85c2-519bce7cc57f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "LENGTH = 1000\n",
    "def crop(signal):\n",
    "    return signal[0:LENGTH, :]\n",
    "\n",
    "def downsample(signal, ratio=5):\n",
    "    # Signal of size (L, 12) to (L/ratio, 12)\n",
    "    return signal[::ratio, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "c91e3c23-4e35-4175-8543-1469adf37a1c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 522/522 [00:15<00:00, 33.98it/s]\n"
     ]
    }
   ],
   "source": [
    "records = []\n",
    "labels = []\n",
    "\n",
    "for key in tqdm(meta['key']):\n",
    "    signal = data[key]\n",
    "    \n",
    "    # Drop the last 3 channels:\n",
    "    signal = signal[:, :12]\n",
    "    assert signal.shape[1] == 12\n",
    "    \n",
    "    # Downsample x10\n",
    "    signal = downsample(signal, 10)\n",
    "    \n",
    "    # Crop\n",
    "    signal = crop(signal)\n",
    "    assert signal.shape == (LENGTH, 12)\n",
    "        \n",
    "    records.append(signal)\n",
    "    \n",
    "    # Get label:\n",
    "    d = meta['Reason_for_admission'][meta['key'] == key].values[0]\n",
    "    labels.append(get_label(d))\n",
    "\n",
    "X, y = np.array(records), np.array(labels)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "ca96ef45-e081-4824-8e27-58b9fb3f142f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1],\n",
       "       [0, 1],\n",
       "       [0, 1],\n",
       "       ...,\n",
       "       [0, 1],\n",
       "       [0, 1],\n",
       "       [0, 1]])"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "8fd740f2-2f16-4fd3-a287-3a2e604bf131",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(522, 1000, 12)\n"
     ]
    }
   ],
   "source": [
    "assert X.shape[0], y.shape[0]\n",
    "    \n",
    "print(X.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "2d3759b3-fda3-4ab9-9eee-67fc178b3beb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "### SPLIT\n",
    "X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=1)\n",
    "\n",
    "assert X_train.shape[1] == 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "34d5d4e1-89b4-4ac4-a690-09d2352f038c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(173, 2)"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "bf2b883c-b8aa-458a-b08a-5cf3dd95874d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.16763006, 0.69364162])"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_test.sum(axis=0)/len(y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "df4aed41-2973-483b-9d00-d569bb7f8b8d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "X_train = X_train.astype(np.float64)\n",
    "X_test = X_test.astype(np.float64)\n",
    "y_train = y_train.astype(np.float64)\n",
    "y_test = y_test.astype(np.float64)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c8c95ec0-0fe1-48de-b5aa-559d057eb4ef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['./data/y_test.joblib']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joblib.dump(X_train, './data/X_train.joblib')\n",
    "joblib.dump(y_train, './data/y_train.joblib')\n",
    "joblib.dump(X_test, './data/X_test.joblib')\n",
    "joblib.dump(y_test, './data/y_test.joblib')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cdff1b4b-d667-4249-8143-fdbd391f45ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "# X_train = joblib.load('./data/X_train.joblib')\n",
    "# X_test = joblib.load('./data/X_test.joblib')\n",
    "# y_train = joblib.load('./data/y_train.joblib')\n",
    "# y_test = joblib.load('./data/y_test.joblib')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a748d649-c807-46fc-8cb3-231d7f3e17ac",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ecg",
   "language": "python",
   "name": "ecg"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
